package LibDL.recommender.data;

import LibDL.data.LibArrayInt;
import LibDL.recommender.data.utils.NumUtils;
import net.librec.data.DataModel;
import net.librec.math.structure.DenseMatrix;
import net.librec.math.structure.SequentialAccessSparseMatrix;
import net.librec.math.structure.SequentialSparseVector;
import net.librec.recommender.item.KeyValue;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.List;

public class Sequence {
    public DenseMatrix user_ids;
    public DenseMatrix sequences;
    public DenseMatrix targets;

    public long L;
    public Long T; // TODO long

    public int num_items;
    public int num_users;

    public LibArrayInt item_ids;
    public DataModel dataModel;

    private Sequence(DenseMatrix user_ids, DenseMatrix sequences, @Nullable DenseMatrix targets,
                     @Nullable DataModel dataModel, int num_users, int num_items) {
        this.user_ids = user_ids;
        this.sequences = sequences;
        this.targets = targets;

        this.L = sequences.columnSize();
        this.T = null;
        if (targets != null && NumUtils.any(targets))
            this.T = (long) targets.columnSize();

        this.dataModel = dataModel;
        this.num_users = num_users;
        this.num_items = num_items;
    }

    public static Sequence[] to_sequence(List<Integer> users, List<Integer> items, boolean isTrain,
                                       int train_num_users, int train_num_items,
                                       int num_users, int num_items,
                                       int sequence_length, int target_length) {
        List<Integer> indices = NumUtils.argsort(users);
        users = NumUtils.index_select(users, indices);
        items = NumUtils.index_select(items, indices);

        List<Integer> unique = new ArrayList<>();
        List<List<Integer>> subLists = new ArrayList<>();
        int temp = users.get(0);
        unique.add(temp);
        int index = 0;
        for (int i = 1; i < users.size(); i++) {
            int k = users.get(i);
            if (users.get(i) != temp) {
                unique.add(k);
                temp = k;
                subLists.add(items.subList(index, i));
                index = i;
            }
        }
        subLists.add(items.subList(index, users.size()));

        int max_sequence_length = sequence_length + target_length;

        int num_subsequences = 0;
        for (List<Integer> subList : subLists) {
            int c = subList.size();
            if (c >= max_sequence_length)
                num_subsequences += c - max_sequence_length + 1;
            else
                num_subsequences += 1;
        }

        DenseMatrix sequences = new DenseMatrix(num_subsequences, sequence_length);
        DenseMatrix sequences_targets = new DenseMatrix(num_subsequences, target_length);
        DenseMatrix sequence_users = new DenseMatrix(1, num_subsequences);

        DenseMatrix test_sequences = new DenseMatrix(num_users, sequence_length);
        DenseMatrix test_users = new DenseMatrix(1, num_users);

        int i = 0;
        int _uid = Integer.MAX_VALUE;

        for (KeyValue<Integer, List<Integer>> pair : _generate_sequences(subLists, max_sequence_length, num_items)) {
            int uid = pair.getKey();
            List<Integer> item_seq = pair.getValue();
            if (uid != _uid) {
                int len = item_seq.size();
                // test_sequences[uid][:] = item_seq[-sequence_length:]
                NumUtils.setItem(test_sequences.row(uid), item_seq.subList(len - sequence_length, len), 0);
                test_users.set(0, uid, uid); // test_users[uid] = uid
                _uid = uid;
            }

            int len = item_seq.size();
            //sequences_targets[i][:] = item_seq[-target_length:]
            NumUtils.setItem(sequences_targets.row(i), item_seq.subList(len - target_length, len), 0);
            //sequences[i][:] = item_seq[:sequence_length]
            NumUtils.setItem(sequences.row(i), item_seq.subList(0, sequence_length), 0);
            //sequence_users[i] = uid;
            sequence_users.set(0, i, uid);
            i++;
        }

        Sequence sequences_ = new Sequence(
                sequence_users, sequences, sequences_targets, null, train_num_users, train_num_items);
        Sequence test_sequences_ = new Sequence(
                test_users, test_sequences, null, null, num_users, num_items);
        return new Sequence[]{sequences_, test_sequences_};

    }

    public static Sequence[] to_sequence(DataModel dataModel, int sequence_length, int target_length) {
//        /* change the item index start from 1 as 0 is used for padding in sequences */
//        item_map.forEach((k, v) -> {
//            item_map.replace(k, v + 1);
//        });
//        item_ids.arrAdd(1);
//        this.num_items += 1;

        int num_users = dataModel.getUserMappingData().size();
        int num_items = dataModel.getItemMappingData().size();

        int max_sequence_length = sequence_length + target_length;

        SequentialAccessSparseMatrix train = (SequentialAccessSparseMatrix) dataModel.getTrainDataSet();

        int num_subsequences = 0;
        for (int i = 0; i < train.rowSize(); i++) {
            int c = train.row(i).size();
            if (c >= max_sequence_length)
                num_subsequences += c - max_sequence_length + 1;
            else
                num_subsequences += 1;
        }

        DenseMatrix sequences = new DenseMatrix(num_subsequences, sequence_length);
        DenseMatrix sequences_targets = new DenseMatrix(num_subsequences, target_length);
        DenseMatrix sequence_users = new DenseMatrix(1, num_subsequences);

        DenseMatrix test_sequences = new DenseMatrix(num_users, sequence_length);
        DenseMatrix test_users = new DenseMatrix(1, num_users);

        int i = 0;
        int _uid = Integer.MAX_VALUE;

        for (KeyValue<Integer, LibArrayInt> pair : _generate_sequences(train, max_sequence_length, num_items)) {
            int uid = pair.getKey();
            LibArrayInt item_seq = pair.getValue();
            if (uid != _uid) {
                int len = item_seq.size();
                // test_sequences[uid][:] = item_seq[-sequence_length:]
                NumUtils.setItem(test_sequences.row(uid), item_seq.subList(len - sequence_length, len), 0);
                test_users.set(0, uid, uid); // test_users[uid] = uid
                _uid = uid;
            }

            int len = item_seq.size();
            //sequences_targets[i][:] = item_seq[-target_length:]
            NumUtils.setItem(sequences_targets.row(i), item_seq.subList(len - target_length, len), 0);
            //sequences[i][:] = item_seq[:sequence_length]
            NumUtils.setItem(sequences.row(i), item_seq.subList(0, sequence_length), 0);
            //sequence_users[i] = uid;
            sequence_users.set(0, i, uid);
            i++;
        }

        Sequence sequences_ = new Sequence(
                sequence_users, sequences, sequences_targets, dataModel, num_users, num_items);
        Sequence test_sequences_ = new Sequence(
                test_users, test_sequences, null, null, num_users, num_items);
        return new Sequence[]{sequences_, test_sequences_};
    }

    private static ArrayList<KeyValue<Integer, LibArrayInt>> _generate_sequences(
            SequentialAccessSparseMatrix train, int max_sequence_length, int num_items) {

        ArrayList<KeyValue<Integer, LibArrayInt>> pairs = new ArrayList<>();

        for (int i = 0; i < train.rowSize(); i++) {

//            int start_idx = indices.get(i);
//            int stop_idx;
//
//            if (i >= indices.size() - 1)
//                stop_idx = item_ids.size();
//            else
//                stop_idx = indices.get(i + 1);

            ArrayList<LibArrayInt> tensorArrayList = _sliding_window(
                    train.row(i), max_sequence_length, 1, num_items); // item_ids[start_idx:stop_idx]
            for (LibArrayInt seq : tensorArrayList) {
                pairs.add(new KeyValue<>(i, seq)); // TODO check `pairs.add(new KeyValue<>(user_ids.get(i), seq));`
            }
        }
        return pairs;
    }

    /**
     * @param num_items Pad sequence with `num_items`
     * */
    private static ArrayList<LibArrayInt> _sliding_window(SequentialSparseVector list, int window_size, int step_size, int num_items) {
        LibArrayInt _list = new LibArrayInt(list.getIndices(), list.size());
        ArrayList<LibArrayInt> tensorArrayList = new ArrayList<>();
        if (_list.size() - window_size >= 0) {
            for (int i = _list.size(); i > 0; i = i - step_size) {
                if (i - window_size >= 0) {
                    tensorArrayList.add(_list.subList(i - window_size, i)); // yield tensor[i - window_size:i]
                } else
                    break;
            }
        } else {
            int num_paddings = window_size - _list.size();
            /* Pad sequence with Number(num_items) if it is shorter than windows size. */
            LibArrayInt left = new LibArrayInt(num_paddings);
            for (int i = 0; i < num_paddings; i++) {
                left.add(0); // TODO CHECK
            }
            left.addAll(_list);
            tensorArrayList.add(left);
        }
        return tensorArrayList;
    }

    private static ArrayList<KeyValue<Integer, List<Integer>>> _generate_sequences(
            List<List<Integer>> train, int max_sequence_length, int num_items) {

        ArrayList<KeyValue<Integer, List<Integer>>> pairs = new ArrayList<>();

        for (int i = 0; i < train.size(); i++) {

//            int start_idx = indices.get(i);
//            int stop_idx;
//
//            if (i >= indices.size() - 1)
//                stop_idx = item_ids.size();
//            else
//                stop_idx = indices.get(i + 1);

            ArrayList<List<Integer>> tensorArrayList = _sliding_window(
                    train.get(i), max_sequence_length, 1, num_items); // item_ids[start_idx:stop_idx]
            for (List<Integer> seq : tensorArrayList) {
                pairs.add(new KeyValue<>(i, seq)); // TODO check `pairs.add(new KeyValue<>(user_ids.get(i), seq));`
            }
        }
        return pairs;
    }

    /**
     * @param num_items Pad sequence with `num_items`
     * */
    private static ArrayList<List<Integer>> _sliding_window(List<Integer> list, int window_size, int step_size, int num_items) {
        ArrayList<List<Integer>> tensorArrayList = new ArrayList<>();
        if (list.size() - window_size >= 0) {
            for (int i = list.size(); i > 0; i = i - step_size) {
                if (i - window_size >= 0) {
                    tensorArrayList.add(list.subList(i - window_size, i)); // yield tensor[i - window_size:i]
                } else
                    break;
            }
        } else {
            int num_paddings = window_size - list.size();
            /* Pad sequence with Number(num_items) if it is shorter than windows size. */
            List<Integer> left = new ArrayList<>(num_paddings);
            for (int i = 0; i < num_paddings; i++) {
                left.add(0); // TODO CHECK
            }
            left.addAll(list);
            tensorArrayList.add(left);
        }
        return tensorArrayList;
    }
}
